from rest_framework.response import Response
import rest_framework

class ApiResponse(Response):
    """
    统一的API响应格式
    """
    def __init__(self, code=200, message="success", data=None, status=None, headers=None, **kwargs):
        body = {
            'code': code,
            'message': message,
            'data': data
        }
        # 添加额外的字段
        body.update(kwargs)

        # 如果没有指定status，则根据code设置
        if status is None:
            if code >= 400:
                status = rest_framework.status.HTTP_400_BAD_REQUEST
            else:
                status = rest_framework.status.HTTP_200_OK
                

        super().__init__(data=body, status=status, headers=headers)

def success(data=None, message="success", **kwargs):
    """成功响应"""
    return ApiResponse(code=200, message=message, data=data, **kwargs)

def error(message="error", code=400, data=None, **kwargs):
    """错误响应"""
    return ApiResponse(code=code, message=message, data=data, **kwargs)